import logging
import random
import pandas as pd
from openpyxl import load_workbook
from openpyxl.styles import Alignment
import os
from datetime import datetime, timedelta

# 配置logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 地区代码示例，这里只用北京的一个区作为例子
area_codes = ['110101']  # 北京市东城区

# 权重因子
weights = [7, 9, 10, 5, 8, 4, 2, 1, 6, 3, 7, 9, 10, 5, 8, 4, 2]

# 校验码表
check_codes = '10X98765432'

# 常见姓氏列表（可以扩展）
surnames = [
    "元", "卜", "顾", "孟", "平", "黄", "和", "穆", "萧", "尹",
    "李", "王", "张", "刘", "陈", "杨", "赵", "周", "吴", "徐",
    "孙", "胡", "朱", "高", "林", "何", "郭", "马", "罗", "梁",
    "宋", "郑", "谢", "韩", "唐", "冯", "于", "董", "肖", "程",
    "曹", "袁", "邓", "许", "傅", "沈", "曾", "彭", "吕", "苏",
    "卢", "蒋", "蔡", "贾", "丁", "魏", "薛", "叶", "阎", "余",
    "潘", "杜", "戴", "夏", "钟", "汪", "田", "任", "姜", "范",
    "方", "石", "姚", "谭", "廖", "邹", "熊", "金", "陆", "郝",
    "孔", "白", "崔", "康", "毛", "邱", "秦", "江", "史", "侯",
    "邵", "孟", "龙", "万", "段", "雷", "钱", "汤", "黎", "易"
]

first_names = [
    "红", "泽", "香", "猛", "霞", "骏", "桂芬", "河", "娜", "浩",
    "伟", "芳", "娜", "强", "静", "明", "杰", "敏", "超", "秀",
    "俊", "勇", "涛", "鹏", "辉", "军", "艳", "玲", "雪", "婷",
    "燕", "波", "洋", "飞", "亮", "松", "刚", "峰", "健", "宁",
    "丹", "梅", "兰", "菊", "华", "琴", "珍", "莉", "莎", "蓉",
    "怡", "欣", "悦", "琪", "瑶", "珊", "妍", "露", "涵", "蕾",
    "凯", "翔", "宇", "轩", "哲", "博", "瑞", "安", "乐", "嘉",
    "豪", "毅", "彬", "锋", "磊", "晨", "曦", "昊", "霖", "森",
    "楠", "柏", "桦", "枫", "柳", "桃", "杏", "莲", "荷", "芝",
    "薇", "蔓", "菁", "菲", "茵", "茜", "茹", "萱", "蓓", "蕊"
]

def load_occupation_data(filename='zhiYeBianMa.xlsx'):
    """从Excel文件中加载职业类别代码、名称和职业等级"""
    df = pd.read_excel(filename)
    occupation_map = dict(zip(df['职业类别代码'], df['职业类别名称']))
    occupation_level_map = dict(zip(df['职业类别代码'], df['职业等级']))  # 新增：加载职业等级
    return occupation_map, occupation_level_map

def load_employer_data(filename='company_name.xlsx'):
    """从Excel文件中加载实际用工单位名称"""
    df = pd.read_excel(filename)
    employers = df['实际用工单位名称'].dropna().tolist()
    return employers

def load_location_data(filename='locations.xlsx'):
    """从Excel文件中加载雇员工作地点（省）和（市）"""
    df = pd.read_excel(filename)
    locations = df[['雇员工作地点（省）', '雇员工作地点（市）']].dropna().values.tolist()
    return locations

def generate_unique_id_number(existing_ids, base_year):
    while True:
        id_number = _generate_id_number(base_year)
        if id_number not in existing_ids:
            return id_number

def _generate_id_number(base_year):
    current_date = datetime.now()
    one_year_ago = current_date - timedelta(days=365)

    birth_year = str(base_year)
    birth_month = f"{random.randint(1, 12):02d}"
    birth_day = f"{random.randint(1, 28):02d}"  # 简化处理，不考虑每个月的具体天数
    birth_date_str = f"{birth_year}{birth_month}{birth_day}"

    birth_date = datetime.strptime(birth_date_str, '%Y%m%d')
    if (current_date - birth_date).days < 365:
        raise ValueError("生成的身份证号码对应的出生日期必须至少为一年以前")

    sequence_code = f"{random.randint(1, 999):03d}"
    id_without_check = random.choice(area_codes) + birth_date_str + sequence_code
    check_sum = sum([int(id_without_check[i]) * weights[i] for i in range(17)])
    check_code = check_codes[check_sum % 11]
    id_number = id_without_check + check_code

    if len(id_number) != 18 or not (
            id_number[:-1].isdigit() and (id_number[-1].isdigit() or id_number[-1].upper() == 'X')):
        raise ValueError(f"生成的身份证号码 {id_number} 不符合18位的要求")

    return id_number

def extract_birth_date(id_number):
    if len(id_number) != 18 or not (
            id_number[:-1].isdigit() and (id_number[-1].isdigit() or id_number[-1].upper() == 'X')):
        raise ValueError(f"身份证号码 {id_number} 必须是18位数字或包含最后一位为X")

    birth_date_str = id_number[6:14]
    try:
        birth_date = pd.to_datetime(birth_date_str, format='%Y%m%d')
        return birth_date.strftime('%Y-%m-%d')
    except ValueError:
        raise ValueError(f"身份证号码中的出生日期部分 {birth_date_str} 不是有效日期")

def calculate_age(birth_date_str):
    """根据身份证号码中的出生日期计算年龄"""
    birth_date = datetime.strptime(birth_date_str, '%Y%m%d')
    today = datetime.today()
    age = today.year - birth_date.year - ((today.month, today.day) < (birth_date.month, birth_date.day))
    return age

def generate_random_name():
    surname = random.choice(surnames)
    first_name_length = random.randint(1, 2)
    first_name = ''.join(random.choices(first_names, k=first_name_length))
    return f"{surname}{first_name}"

def create_dummy_data(count=5, age_multiplier=1, age_step=1, occupation_map=None, occupation_level_map=None, employer_list=None, location_list=None):
    data = []
    existing_ids = set()

    # 设置初始基准年份
    base_year = 1980

    # 计算需要生成的组数
    group_count = count // age_multiplier
    remainder = count % age_multiplier

    for group in range(group_count):
        # 根据年龄步进和当前组号计算出生年份
        birth_year = base_year + (group * age_step)
        birth_date_str = f"{birth_year}0101"
        age = calculate_age(birth_date_str)

        for _ in range(age_multiplier):
            try:
                id_number = generate_unique_id_number(existing_ids, birth_year)
                existing_ids.add(id_number)
                gender = '男' if int(id_number[16]) % 2 else '女'
                name = generate_random_name()

                # 随机选择职业
                if occupation_map and occupation_level_map:
                    occupation_code, occupation_name = random.choice(list(occupation_map.items()))
                    occupation_level = occupation_level_map[occupation_code]  # 获取职业等级
                else:
                    occupation_code, occupation_name, occupation_level = "未知代码", "未知名称", "未知等级"

                # 随机选择实际用工单位名称
                employer = random.choice(employer_list) if employer_list else "未知单位"

                # 随机选择工作地点
                if location_list:
                    province, city = random.choice(location_list)
                else:
                    province, city = "未知省份", "未知城市"

                entry = {
                    "序号": len(data) + 1,
                    "姓名": name,
                    "证件类型": "居民身份证",
                    "证件号码": id_number,
                    "职业类别代码": occupation_code,
                    "职业类别名称": occupation_name,
                    "性别": gender,
                    "出生日期": extract_birth_date(id_number),
                    "实际用工单位名称": employer,
                    "雇员工作地点（省）": province,
                    "雇员工作地点（市）": city,
                    "雇员年龄": age,  # 新增字段：雇员年龄
                    "职业等级": occupation_level  # 新增字段：职业等级
                }
                data.append(entry)

            except ValueError as e:
                logging.error(f"生成第 {len(data) + 1} 条记录时出错: {e}")
                continue

    # 处理剩余的雇员
    for _ in range(remainder):
        try:
            id_number = generate_unique_id_number(existing_ids, base_year + (group_count * age_step))
            existing_ids.add(id_number)
            gender = '男' if int(id_number[16]) % 2 else '女'
            name = generate_random_name()

            # 计算年龄
            birth_date_str = id_number[6:14]
            age = calculate_age(birth_date_str)

            # 随机选择职业
            if occupation_map and occupation_level_map:
                occupation_code, occupation_name = random.choice(list(occupation_map.items()))
                occupation_level = occupation_level_map[occupation_code]  # 获取职业等级
            else:
                occupation_code, occupation_name, occupation_level = "未知代码", "未知名称", "未知等级"

            # 随机选择实际用工单位名称
            employer = random.choice(employer_list) if employer_list else "未知单位"

            # 随机选择工作地点
            if location_list:
                province, city = random.choice(location_list)
            else:
                province, city = "未知省份", "未知城市"

            entry = {
                "序号": len(data) + 1,
                "姓名": name,
                "证件类型": "居民身份证",
                "证件号码": id_number,
                "职业类别代码": occupation_code,
                "职业类别名称": occupation_name,
                "性别": gender,
                "出生日期": extract_birth_date(id_number),
                "实际用工单位名称": employer,
                "雇员工作地点（省）": province,
                "雇员工作地点（市）": city,
                "雇员年龄": age,  # 新增字段：雇员年龄
                "职业等级": occupation_level  # 新增字段：职业等级
            }
            data.append(entry)

        except ValueError as e:
            logging.error(f"生成第 {len(data) + 1} 条记录时出错: {e}")
            continue

        if len(data) % 1000 == 0:
            logging.info(f"已成功生成 {len(data)} 条记录...")

    return data

def fill_template_with_data(template_filename, output_filename, data, start_row=3):
    wb = load_workbook(template_filename)
    ws = wb.active

    headers = [cell.value for cell in ws[2]]

    for row_num, entry in enumerate(data, start=start_row):
        for col_num, header in enumerate(headers, start=1):
            cell = ws.cell(row=row_num, column=col_num)
            cell.value = str(entry.get(header, ""))

            if header == '证件号码':
                cell.number_format = '@'
            elif header == '出生日期':
                cell.number_format = 'YYYY-MM-DD'

            cell.alignment = Alignment(horizontal="center", vertical="center")

    # 修改为指定路径 D:\software\chrome
    output_path = r"D:\software\chrome"  # 使用原始字符串避免转义问题
    excel_filename = os.path.join(output_path, output_filename)

    logging.info(f"即将保存文件到: {excel_filename}")

    try:
        wb.save(excel_filename)
        logging.info(f"数据已成功写入 {excel_filename}")
    except PermissionError:
        logging.error(f"无法保存文件 {excel_filename}，请确保文件未被其他程序占用或您有权限写入该位置。")
        exit(1)
    except Exception as e:
        logging.error(f"保存文件时发生错误: {e}")
        exit(1)
def generate_timestamp_filename(prefix="output"):
    timestamp = datetime.now().strftime("%H%M")
    return f"{prefix}_{timestamp}.xlsx"

if __name__ == "__main__":
    template_filename = 'template.xlsx'

    if not os.path.isfile(template_filename):
        logging.error(f"模板文件 {template_filename} 不存在，请检查路径。")
        exit(1)

    # 加载职业类别数据、用工单位数据和工作地点数据
    occupation_map, occupation_level_map = load_occupation_data('zhiYeBianMa.xlsx')
    employer_list = load_employer_data('company_name.xlsx')
    location_list = load_location_data('locations.xlsx')

    # 设置年龄倍数开关和年龄步进
    age_multiplier = 10
    age_step = 5  # 每个雇员之间的年龄差为2岁
    dummy_data = create_dummy_data(
        count=5,  # 你可以根据需要调整数量
        age_multiplier=age_multiplier,
        age_step=age_step,
        occupation_map=occupation_map,
        occupation_level_map=occupation_level_map,
        employer_list=employer_list,
        location_list=location_list
    )

    output_filename = generate_timestamp_filename(prefix="")
    fill_template_with_data(template_filename, output_filename, dummy_data, start_row=3)

    logging.info(f"已成功创建文件: {output_filename}")